import torch
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data import DataLoader
from dataset.vocab import WordVocab
from model.LayoutLMv2Model import LayoutLMv2Model
from trainer.integrated_model import LAYOUYLMV2
from trainer.optim_schedule import ScheduledOptim

import tqdm
class LayoutLMv2Trainer:


    def __init__(self, layoutlm: LayoutLMv2Model, config,
                 train_dataloader: DataLoader, test_dataloader: DataLoader = None,
                 lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.01, warmup_steps=10000,
                 with_cuda: bool = True, cuda_devices=None, log_freq: int = 10,total_acc: float = 0):
        """
        :param layoutlm: LayoutLMv2Model model which you want to train
        :param vocab_size: total word vocab size
        :param train_dataloader: train dataset data loader
        :param test_dataloader: test dataset data loader [can be None]
        :param lr: learning rate of optimizer
        :param betas: Adam optimizer betas
        :param weight_decay: Adam optimizer weight decay param
        :param with_cuda: traning with cuda
        :param log_freq: logging frequency of the batch iteration
        """

        # Setup cuda device for layoutlm training, argument -c, --cuda should be true
        cuda_condition = torch.cuda.is_available() and with_cuda
        self.device = torch.device("cuda:0" if cuda_condition else "cpu")

        # This layoutlm model will be saved every epoch
        self.layoutlm = layoutlm
        self.config = config
        # Initialize the layoutlm Language Model, with LAYOUYLMV2 model
        self.model = LAYOUYLMV2(self.layoutlm, self.config).to(self.device)

        # Distributed GPU training if CUDA can detect more than 1 GPU
        if with_cuda and torch.cuda.device_count() > 1:
            print("Using %d GPUS for layoutlm" % torch.cuda.device_count())
            self.model = nn.DataParallel(self.model, device_ids=cuda_devices)

        # Setting the train and test data loader
        self.train_data = train_dataloader
        self.test_data = test_dataloader

        # Setting the Adam optimizer with hyper-param
        self.optim = Adam(self.model.parameters(), lr=lr, betas=betas, weight_decay=weight_decay)
        self.optim_schedule = ScheduledOptim(self.optim, self.config.hidden_size, n_warmup_steps=config.TRAIN.WARMUP_LR)

        # Using Negative Log Likelihood Loss function for predicting the masked_token
        self.criterion = nn.NLLLoss(ignore_index=0)

        self.log_freq = log_freq
        self.total_acc = total_acc
        print("Total Parameters:", sum([p.nelement() for p in self.model.parameters()]))

    def train(self, epoch):
        self.iteration(epoch, self.train_data)

    def test(self, epoch):
        self.iteration(epoch, self.test_data, train=False)

    def iteration(self, epoch, data_loader, train=True):
        """
        loop over the data_loader for training or testing
        if on train status, backward operation is activated
        and also auto save the model every peoch

        :param epoch: current epoch index
        :param data_loader: torch.utils.data.DataLoader for iteration
        :param train: boolean value of is train or test
        :return: None
        """
        str_code = "train" if train else "test"

        # Setting the tqdm progress bar
        data_iter = tqdm.tqdm(enumerate(data_loader),
                              desc="EP_%s:%d" % (str_code, epoch),
                              total=len(data_loader),
                              bar_format="{l_bar}{r_bar}")

        avg_loss = 0.0
        total_correct = 0
        total_element = 0
        total_acc = 0
        for i, data in data_iter:
            # 0. batch_data will be sent into the device(GPU or cpu)
            data = {key: value.to(self.device) for key, value in data.items()}

            # 1. forward the covered_prediction and masked_lm model
            covered_output, mask_lm_output = self.model(data["curpus_input"],data["box_input"], data["img_input"])


            # 2-1. NLL(negative log likelihood) loss of is_cover classification result
            covered_loss = self.criterion(covered_output.transpose(1,2), data["covered_label"])

            # 2-2. NLLLoss of predicting masked token word
            mask_loss = self.criterion(mask_lm_output.transpose(1, 2), data["curpus_label"])

            # 2-3. Adding covered_loss and mask_loss : 3.4 Pre-training Procedure
            loss = covered_loss + mask_loss

            # 3. backward and optimization only in train
            if train:
                self.optim_schedule.zero_grad()
                loss.backward()
                self.optim_schedule.step_and_update_lr()

            # next sentence prediction accuracy
            correct = covered_output.argmax(dim=-1).eq(data["covered_label"]).sum().item()

            avg_loss += loss.item()
            total_correct += correct
            total_element += data["covered_label"].nelement()

            post_fix = {
                "epoch": epoch,
                "iter": i,
                "avg_loss": avg_loss / (i + 1),
                "avg_acc": total_correct / total_element,
                "loss": loss.item()
            }
            print(post_fix)
            if i % self.log_freq == 0:
                data_iter.write(str(post_fix))
        self.total_acc =  total_correct  / total_element

        print("EP%d_%s, avg_loss=" % (epoch, str_code), avg_loss / len(data_iter), "total_acc=",
              total_correct / total_element)


    def save(self, epoch, file_path="lmv2_trained.model"):
        """
        Saving the current layoutlm model on file_path

        :param epoch: current epoch number
        :param file_path: model output path which gonna be file_path+"ep%d" % epoch
        :return: final_output_path
        """

        output_path = file_path + "lmv2_trained.model%d" % epoch
        torch.save(self.model.state_dict(), output_path)
        self.model.to(self.device)
        print("EP:%d Model Saved on:" % epoch, output_path)
        return output_path

if __name__=="__main__":
    from dataset.build_data import dataloder
    from configs.config import _C as config
    train_data = dataloder(config)

    trainer = LayoutLMv2Trainer(LayoutLMv2Model,config,train_data)
    for epoch in range(config.TRAIN.EPOCHS):
        trainer.train(epoch)
        acc = 0
        last_acc = trainer.total_acc
        print('last_acc',last_acc)
        # if acc <last_acc:
        #     acc = last_acc
        #     trainer.save(epoch, config.output_path)

